import sys
from PyQt6.QtWidgets import QApplication, QWidget, QLabel, QLineEdit, QPushButton, QTextEdit, QComboBox, QFileDialog
from netmiko import ConnectHandler
from datetime import datetime

class InspectionApp(QWidget):
    def __init__(self):
        super().__init__()

        self.initUI()

    def initUI(self):
        # 创建标签和输入框
        self.label_ip = QLabel('IP地址:', self)
        self.label_ip.move(20, 20)
        self.input_ip = QLineEdit(self)
        self.input_ip.move(100, 20)

        self.label_port = QLabel('端口号:', self)
        self.label_port.move(20, 60)
        self.input_port = QLineEdit(self)
        self.input_port.move(100, 60)

        self.label_username = QLabel('用户名:', self)
        self.label_username.move(20, 100)
        self.input_username = QLineEdit(self)
        self.input_username.move(100, 100)

        self.label_password = QLabel('密码:', self)
        self.label_password.move(20, 140)
        self.input_password = QLineEdit(self)
        self.input_password.setEchoMode(QLineEdit.EchoMode.Password)
        self.input_password.move(100, 140)

        self.label_device_type = QLabel('设备类型:', self,)
        self.label_device_type.move(20, 180)
        self.combo_device_type = QComboBox(self)
        self.combo_device_type.addItems(['cisco_ios', 'huawei', 'arista_eos', 'juniper','hp_comware'])
        self.combo_device_type.move(100, 180)

        self.label_tips = QLabel('tips:华三交换机使用hp_comware',self)
        self.label_tips.move(20, 210)

        # 创建按钮
        self.btn_inspect = QPushButton('开始巡检', self)
        self.btn_inspect.move(20, 240)
        self.btn_inspect.clicked.connect(self.inspect)

        self.btn_select_dir = QPushButton('选择保存目录', self)
        self.btn_select_dir.move(150, 240)
        self.btn_select_dir.clicked.connect(self.select_dir)

        # 创建日志显示框
        self.log_display = QTextEdit(self)
        self.log_display.setReadOnly(True)
        self.log_display.setGeometry(250, 20, 300, 240)

        # 设置窗口大小和标题
        self.setGeometry(200, 200, 580, 300)
        self.setWindowTitle('交换机巡检工具_20231108_V1_by.仓鼠')

    def inspect(self):
        # 获取用户输入的信息
        ip = self.input_ip.text()
        port = self.input_port.text()
        username = self.input_username.text()
        password = self.input_password.text()
        device_type = self.combo_device_type.currentText()

        # 连接交换机
        device = {
            'device_type': device_type,
            'ip': ip,
            'port': port,
            'username': username,
            'password': password,
        }

        try:
            connection = ConnectHandler(**device)
            connection.enable()

            # 执行巡检命令
            commands = {
                'cisco_ios': [
                    'show version',
                    'show interfaces status',
                    'show ip interface brief',
                ],
                'huawei': [
                    'display version',
                    'display interface brief',
                    'display ip interface brief',
                ],
                'arista_eos': [
                    'show version',
                    'show interfaces status',
                    'show ip interface brief',
                ],
                'juniper': [
                    'show version',
                    'show interfaces terse',
                    'show route summary',
                ],
                'hp_comware': [
                    'display version',
                    'display interface brief',
                    'display ip interface brief',
                ],
            }

            log = ''
            for command in commands.get(device_type, []):
                output = connection.send_command(command)
                log += f"Output for command '{command}':\n{output}\n"

            # 将结果显示在日志框中
            self.log_display.setPlainText(log)

            # 保存到文件
            now = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
            filename = f"{now}_{ip}.txt"

            if hasattr(self, 'selected_dir'):
                filename = f"{self.selected_dir}/{filename}"

            with open(filename, 'w') as f:
                f.write(log)

            # 关闭连接
            connection.disconnect()
        except Exception as e:
            self.log_display.setPlainText(f"Error: {str(e)}")

    def select_dir(self):
        dialog = QFileDialog()
        dialog.setFileMode(QFileDialog.FileMode.Directory)
        dialog.setViewMode(QFileDialog.ViewMode.List)
        dialog.setOptions(QFileDialog.Option.ShowDirsOnly)
        self.selected_dir = dialog.getExistingDirectory(self, "选择保存目录", "")
        if not self.selected_dir:
            self.selected_dir = None

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = InspectionApp()
    ex.show()
    sys.exit(app.exec())
